# Author : ZZH
# Date : 2025/2/6
"""
 视频分帧处理
"""
import cv2
import os


def extract_key_frames(video_path, output_dir, interval_sec=2):
    """
    按固定时间间隔提取关键帧
    :param
    video_path: 视频路径
    output_dir: 输出目录
    interval_sec: 每间隔N秒取一帧
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 读取视频
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)             # 获取视频帧率
    frame_interval = int(fps * interval_sec)    # 计算每间隔N秒取一帧
    frame_count = 0

    while cap.isOpened():                       # 用于检查视频捕获对象 cap 是否仍然打开
        ret, frame = cap.read()                 # ret: 是否读取到帧，frame: 读取到的图像帧
        if not ret:
            break
        # 按间隔保存帧到指定路径
        if frame_count % frame_interval == 0:
            output_path = os.path.join(output_dir, f"frame_{frame_count:04d}.jpg")
            cv2.imwrite(output_path, frame)
        frame_count += 1

    # 释放资源
    cap.release()
    print(f"Saved {frame_count // frame_interval} frames to {output_dir}")


# 每2秒提取一帧
extract_key_frames("test_video.mp4", "extracted_frames", interval_sec=2)